#pragma once
#include <iostream>
using namespace std;
namespace xzq
{
    template <class K, class V>
    struct AVLtreeNode
    {
        AVLtreeNode<K, V> *_left;
        AVLtreeNode<K, V> *_right;
        AVLtreeNode<K, V> *_parent;
        int _bf;
        pair<K, V> _kv;
        AVLtreeNode(const pair<K, V> &kv)
            : _left(nullptr), _right(nullptr), _parent(nullptr), _kv(kv), _bf(0)
        {
        }
    };

    template <class K, class V>
    class AVLtree
    {
        typedef AVLtreeNode<K, V> Node;

    public:
        bool Insert(const pair<K, V> &kv)
        {
            if (_root == nullptr)
            {
                _root = new Node(kv);
                return true;
            }
            Node *parent = nullptr;
            Node *cur = _root;
            while (cur)
            {
                if (cur->_kv.first > kv.first)
                {
                    parent = cur;
                    cur = cur->_left;
                }
                else if (cur->_kv.first < kv.first)
                {
                    parent = cur;
                    cur = cur->_right;
                }
                else
                {
                    return false;
                }
            }
            cur = new Node(kv);
            if (parent->_kv.first < kv.first)
            {
                parent->_right = cur;
                cur->_parent = parent;
            }
            else
            {
                parent->_left = cur;
                cur->_parent = parent;
            }
            while (parent)
            {
                if (cur == parent->_right)
                {
                    parent->_bf++;
                }
                else
                {
                    parent->_bf--;
                }
                if (parent->_bf == 0)
                {
                    break;
                }
                else if (parent->_bf == 1 || parent->_bf == -1)
                {
                    cur = parent;
                    parent = parent->_parent;
                }
                else if (parent->_bf == 2 || parent->_bf == -2)
                {
                    if (parent->_bf == 2)
                    {
                        if (cur->_bf == 1)
                        {
                            RotateL(parent);
                        }
                        else if (cur->_bf == -1)
                        {
                            RotateRL(parent);
                        }
                    }
                    else if (parent->_bf == -2)
                    {
                        if (cur->_bf == -1)
                        {
                            RotateR(parent);
                        }
                        else if (cur->_bf == 1)
                        {
                            RotateLR(parent);
                        }
                    }
                    break;
                }
            }
            return true;
        }
        void RotateL(Node *parent)
        {
            Node *subR = parent->_right;
            Node *subRL = subR->_left;
            parent->_right = subRL;
            if (subRL)
            {
                subRL->_parent = parent;
            }
            subR->_left = parent;
            Node *ppNode = parent->_parent;
            parent->_parent = subR;
            if (_root == parent)
            {
                _root = subR;
                subR->_parent = nullptr;
            }
            else
            {
                if (ppNode->_left == parent)
                {
                    ppNode->_left = subR;
                }
                else
                {
                    ppNode->_right = subR;
                }
                subR->_parent = ppNode;
            }
            parent->_bf = subR->_bf = 0;
        }
        void RotateR(Node *parent)
        {
            Node *subL = parent->_left;
            Node *subLR = subL->_right;
            parent->_left = subLR;
            if (subLR)
            {
                subLR->_parent = parent;
            }
            subL->_right = parent;
            Node *ppNode = parent->_parent;
            parent->_parent = subL;
            if (_root == parent)
            {
                _root = subL;
                subL->_parent = nullptr;
            }
            else
            {
                if (ppNode->_left == parent)
                {
                    ppNode->_left = subL;
                }
                else
                {
                    ppNode->_right = subL;
                }
                subL->_parent = ppNode;
            }
            parent->_bf = subL->_bf = 0;
        }
        void RotateRL(Node *parent)
        {
            Node *subR = parent->_right;
            Node *subRL = subR->_left;
            int bf = subRL->_bf;
            RotateR(parent->_right);
            RotateL(parent);
            if (bf == -1)
            {
                parent->_bf = 0;
                subR->_bf = 1;
                subRL->_bf = 0;
            }
            else if (bf == 1)
            {
                parent->_bf = -1;
                subR->_bf = 0;
                subRL->_bf = 0;
            }
            else if (bf == 0)
            {
                parent->_bf = 0;
                subR->_bf = 0;
                subRL->_bf = 0;
            }
        }
        void RotateLR(Node *parent)
        {
            Node *subL = parent->_left;
            Node *subLR = subL->_right;
            int bf = subLR->_bf;
            RotateL(parent->_left);
            RotateR(parent);
            if (bf == -1)
            {
                subL->_bf = 0;
                parent->_bf = 1;
                subLR->_bf = 0;
            }
            else if (bf == 1)
            {
                subL->_bf = -1;
                parent->_bf = 0;
                subLR->_bf = 0;
            }
            else if (bf == 0)
            {
                subL->_bf = 0;
                parent->_bf = 0;
                subLR->_bf = 0;
            }
        }
        void _Inorder(Node *root)
        {
            if (root == nullptr)
                return;
            _Inorder(root->_left);
            cout << root->_kv.first << " : " << root->_kv.second << " : " << root->_bf << endl;
            _Inorder(root->_right);
        }
        void Inorder()
        {
            _Inorder(_root);
        }
        int Height(Node *root)
        {
            if (root == nullptr)
                return 0;
            int LeftHeight = Height(root->_left);
            int RightHeight = Height(root->_right);
            return LeftHeight > RightHeight ? LeftHeight + 1 : RightHeight + 1;
        }
        bool _IsBalance(Node *root)
        {
            if (root == nullptr)
                return true;
            int LeftHeight = Height(root->_left);
            int RightHeight = Height(root->_right);
            return abs(LeftHeight - RightHeight) < 2 && _IsBalance(root->_left) && _IsBalance(root->_right);
        }
        bool IsBalance()
        {
            return _IsBalance(_root);
        }

    private:
        Node *_root = nullptr;
    };
}